#include "cmd-undocumented.h"

#include "argument.h"
#include "bookmark.h"
#include "comment.h"
#include "console.h"
#include "debugger.h"
#include "function.h"
#include "label.h"
#include "loop.h"
#include "mnemonichelp.h"
#include "symbolinfo.h"
#include "value.h"
#include "variable.h"
#include "zydis_wrapper.h"

bool cbBadCmd(int argc, char* argv[]) {
  duint value = 0;
  int valsize = 0;
  bool isvar = false;
  bool hexonly = false;
  if (valfromstring(*argv, &value, false, false, &valsize, &isvar, &hexonly,
                    true))  // dump variable/value/register/etc
  {
    varset("$ans", value, true);
    if (valsize)
      valsize *= 2;
    else
      valsize = 1;
    char format_str[deflen] = "";
    auto symbolic = SymGetSymbolicName(value);
    if (symbolic.length()) symbolic = " " + symbolic;
    if (isvar)  // and *cmd!='.' and *cmd!='x') //prevent stupid 0=0 stuff
    {
      if (value > 9 && !hexonly) {
        if (!valuesignedcalc())  // signed numbers
#ifdef _WIN64
          sprintf_s(format_str, "%%s=%%.%dllX (%%llud)%%s\n",
                    valsize);  // TODO: This and the following statements use
                               // "%llX" for a "int"-typed variable. Maybe we
                               // can use "%X" everywhere?
#else  // x86
          sprintf_s(format_str, "%%s=%%.%dX (%%ud)%%s\n", valsize);
#endif  //_WIN64
        else
#ifdef _WIN64
          sprintf_s(format_str, "%%s=%%.%dllX (%%lld)%%s\n", valsize);
#else  // x86
          sprintf_s(format_str, "%%s=%%.%dX (%%d)%%s\n", valsize);
#endif  //_WIN64
        dprintf_untranslated(format_str, *argv, value, value, symbolic.c_str());
      } else {
        sprintf_s(format_str, "%%s=%%.%dX%%s\n", valsize);
        dprintf_untranslated(format_str, *argv, value, symbolic.c_str());
      }
    } else {
      if (value > 9 && !hexonly) {
        if (!valuesignedcalc())  // signed numbers
#ifdef _WIN64
          sprintf_s(format_str, "%%s=%%.%dllX (%%llud)%%s\n", valsize);
#else  // x86
          sprintf_s(format_str, "%%s=%%.%dX (%%ud)%%s\n", valsize);
#endif  //_WIN64
        else
#ifdef _WIN64
          sprintf_s(format_str, "%%s=%%.%dllX (%%lld)%%s\n", valsize);
#else  // x86
          sprintf_s(format_str, "%%s=%%.%dX (%%d)%%s\n", valsize);
#endif  //_WIN64
#ifdef _WIN64
        sprintf_s(format_str, "%%.%dllX (%%llud)%%s\n", valsize);
#else  // x86
        sprintf_s(format_str, "%%.%dX (%%ud)%%s\n", valsize);
#endif  //_WIN64
        dprintf_untranslated(format_str, value, value, symbolic.c_str());
      } else {
#ifdef _WIN64
        sprintf_s(format_str, "%%.%dllX%%s\n", valsize);
#else  // x86
        sprintf_s(format_str, "%%.%dX%%s\n", valsize);
#endif  //_WIN64
        dprintf_untranslated(format_str, value, symbolic.c_str());
      }
    }
  } else  // unknown command
  {
    dprintf_untranslated("Unknown command/expression: \"%s\"\n", *argv);
    return false;
  }
  return true;
}

bool cbDebugBenchmark(int argc, char* argv[]) {
  duint addr = MemFindBaseAddr(GetContextDataEx(hActiveThread, UE_CIP), 0);
  DWORD ticks = GetTickCount();
  for (duint i = addr; i < addr + 100000; i++) {
    CommentSet(i, "test", false);
    LabelSet(i, "test", false);
    BookmarkSet(i, false);
    FunctionAdd(i, i, false);
    ArgumentAdd(i, i, false);
    LoopAdd(i, i, false);
  }
  dprintf_untranslated("%ums\n", GetTickCount() - ticks);
  return true;
}

bool cbInstrSetstr(int argc, char* argv[]) {
  if (IsArgumentsLessThan(argc, 3)) return false;
  varnew(argv[1], 0, VAR_USER);
  if (!vargettype(argv[1], 0)) {
    dprintf(QT_TRANSLATE_NOOP("DBG", "No such variable \"%s\"!\n"), argv[1]);
    return false;
  }
  if (!varset(argv[1], argv[2], false)) {
    dprintf(QT_TRANSLATE_NOOP("DBG", "Failed to set variable \"%s\"!\n"),
            argv[1]);
    return false;
  }
  cmddirectexec(StringUtils::sprintf("getstr \"%s\"", argv[1]).c_str());
  return true;
}

bool cbInstrGetstr(int argc, char* argv[]) {
  if (IsArgumentsLessThan(argc, 2)) return false;
  VAR_VALUE_TYPE valtype;
  if (!vargettype(argv[1], 0, &valtype)) {
    dprintf(QT_TRANSLATE_NOOP("DBG", "No such variable \"%s\"!\n"), argv[1]);
    return false;
  }
  if (valtype != VAR_STRING) {
    dprintf(QT_TRANSLATE_NOOP("DBG", "Variable \"%s\" is not a string!\n"),
            argv[1]);
    return false;
  }
  int size;
  if (!varget(argv[1], (char*)0, &size, 0) || !size) {
    dprintf(QT_TRANSLATE_NOOP("DBG", "Failed to get variable size \"%s\"!\n"),
            argv[1]);
    return false;
  }
  Memory<char*> string(size + 1, "cbInstrGetstr:string");
  if (!varget(argv[1], string(), &size, 0)) {
    dprintf(QT_TRANSLATE_NOOP("DBG", "Failed to get variable data \"%s\"!\n"),
            argv[1]);
    return false;
  }
  dprintf_untranslated("%s=\"%s\"\n", argv[1], string());
  return true;
}

bool cbInstrCopystr(int argc, char* argv[]) {
  if (IsArgumentsLessThan(argc, 3)) return false;
  VAR_VALUE_TYPE valtype;
  if (!vargettype(argv[2], 0, &valtype)) {
    dprintf(QT_TRANSLATE_NOOP("DBG", "No such variable \"%s\"!\n"), argv[2]);
    return false;
  }
  if (valtype != VAR_STRING) {
    dprintf(QT_TRANSLATE_NOOP("DBG", "Variable \"%s\" is not a string!\n"),
            argv[2]);
    return false;
  }
  int size;
  if (!varget(argv[2], (char*)0, &size, 0) || !size) {
    dprintf(QT_TRANSLATE_NOOP("DBG", "Failed to get variable size \"%s\"!\n"),
            argv[2]);
    return false;
  }
  Memory<char*> string(size + 1, "cbInstrGetstr:string");
  if (!varget(argv[2], string(), &size, 0)) {
    dprintf(QT_TRANSLATE_NOOP("DBG", "Failed to get variable data \"%s\"!\n"),
            argv[2]);
    return false;
  }
  duint addr;
  if (!valfromstring(argv[1], &addr)) {
    dprintf(QT_TRANSLATE_NOOP("DBG", "Invalid address \"%s\"!\n"), argv[1]);
    return false;
  }
  if (!MemPatch(addr, string(), strlen(string()))) {
    dputs(QT_TRANSLATE_NOOP("DBG", "MemPatch failed!"));
    return false;
  }
  dputs(QT_TRANSLATE_NOOP("DBG", "String written!"));
  GuiUpdateAllViews();
  GuiUpdatePatches();
  return true;
}

bool cbInstrZydis(int argc, char* argv[]) {
  if (IsArgumentsLessThan(argc, 2)) return false;

  duint addr = 0;
  if (!valfromstring(argv[1], &addr) || !MemIsValidReadPtr(addr)) {
    dprintf_untranslated("Invalid address \"%s\"\n", argv[1]);
    return false;
  }

  unsigned char data[16];
  if (!MemRead(addr, data, sizeof(data))) {
    dprintf_untranslated("Could not read memory at %p\n", addr);
    return false;
  }

  if (argc > 2)
    if (!valfromstring(argv[2], &addr, false)) return false;

  Zydis cp;
  if (!cp.Disassemble(addr, data)) {
    dputs_untranslated("Failed to disassemble!\n");
    return false;
  }

  auto instr = cp.GetInstr();
  int argcount = instr->operandCount;
  dputs_untranslated(cp.InstructionText(true).c_str());
  dprintf_untranslated("prefix size: %d\n", instr->raw.prefixes.count);
  if (instr->raw.rex.isDecoded)
    dprintf_untranslated(
        "rex.W: %d, rex.R: %d, rex.X: %d, rex.B: %d, rex.data: %02x\n",
        instr->raw.rex.W, instr->raw.rex.R, instr->raw.rex.X, instr->raw.rex.B,
        instr->raw.rex.data[0]);
  dprintf_untranslated("disp.offset: %d, disp.size: %d\n",
                       instr->raw.disp.offset, instr->raw.disp.size);
  dprintf_untranslated("imm[0].offset: %d, imm[0].size: %d\n",
                       instr->raw.imm[0].offset, instr->raw.imm[0].size);
  dprintf_untranslated("imm[1].offset: %d, imm[1].size: %d\n",
                       instr->raw.imm[1].offset, instr->raw.imm[1].size);
  dprintf_untranslated("size: %d, id: %d, opcount: %d\n", cp.Size(), cp.GetId(),
                       instr->operandCount);
  auto rwstr = [](uint8_t action) {
    switch (action) {
      case ZYDIS_OPERAND_ACTION_READ:
      case ZYDIS_OPERAND_ACTION_CONDREAD:
        return "read";
      case ZYDIS_OPERAND_ACTION_WRITE:
      case ZYDIS_OPERAND_ACTION_CONDWRITE:
        return "write";
      case ZYDIS_OPERAND_ACTION_READWRITE:
      case ZYDIS_OPERAND_ACTION_READ_CONDWRITE:
      case ZYDIS_OPERAND_ACTION_CONDREAD_WRITE:
        return "read+write";
      default:
        return "???";
    }
  };
  auto vis = [](uint8_t visibility) {
    switch (visibility) {
      case ZYDIS_OPERAND_VISIBILITY_INVALID:
        return "invalid";
      case ZYDIS_OPERAND_VISIBILITY_EXPLICIT:
        return "explicit";
      case ZYDIS_OPERAND_VISIBILITY_IMPLICIT:
        return "implicit";
      case ZYDIS_OPERAND_VISIBILITY_HIDDEN:
        return "hidden";
      default:
        return "???";
    }
  };
  for (int i = 0; i < argcount; i++) {
    const auto& op = instr->operands[i];
    dprintf("operand %d (size: %d, access: %s, visibility: %s) \"%s\", ", i + 1,
            op.size, rwstr(op.action), vis(op.visibility),
            cp.OperandText(i).c_str());
    switch (op.type) {
      case ZYDIS_OPERAND_TYPE_REGISTER:
        dprintf_untranslated("register: %s\n", cp.RegName(op.reg.value));
        break;
      case ZYDIS_OPERAND_TYPE_IMMEDIATE:
        dprintf_untranslated("immediate: 0x%p\n", op.imm.value.u);
        break;
      case ZYDIS_OPERAND_TYPE_MEMORY: {
        //[base + index * scale +/- disp]
        const auto& mem = op.mem;
        dprintf_untranslated(
            "memory segment: %s, base: %s, index: %s, scale: %d, displacement: "
            "0x%p\n",
            cp.RegName(mem.segment), cp.RegName(mem.base),
            cp.RegName(mem.index), mem.scale, mem.disp.value);
      } break;
      case ZYDIS_OPERAND_TYPE_POINTER:
        dprintf_untranslated("pointer: %X:%p\n", op.ptr.segment, op.ptr.offset);
        break;
    }
  }

  return true;
}

bool cbInstrVisualize(int argc, char* argv[]) {
  if (IsArgumentsLessThan(argc, 3)) return false;
  duint start;
  duint maxaddr;
  if (!valfromstring(argv[1], &start) || !valfromstring(argv[2], &maxaddr)) {
    dputs_untranslated("Invalid arguments!");
    return false;
  }
  // actual algorithm
  // make sure to set these options in the INI (rest the default theme of
  // x64dbg): DisassemblyBookmarkBackgroundColor = #00FFFF
  // DisassemblyBookmarkColor = #000000
  // DisassemblyHardwareBreakpointBackgroundColor = #00FF00
  // DisassemblyHardwareBreakpointColor = #000000
  // DisassemblyBreakpointBackgroundColor = #FF0000
  // DisassemblyBreakpointColor = #000000
  {
    // initialize
    Zydis _cp;
    duint _base = start;
    duint _size = maxaddr - start;
    Memory<unsigned char*> _data(_size);
    MemRead(_base, _data(), _size);
    FunctionClear();

    // linear search with some trickery
    duint end = 0;
    duint jumpback = 0;
    for (duint addr = start, fardest = 0; addr < maxaddr;) {
      // update GUI
      BpClear();
      BookmarkClear();
      LabelClear();
      SetContextDataEx(fdProcessInfo->hThread, UE_CIP, addr);
      if (end) BpNew(end, true, false, 0, BPNORMAL, 0, nullptr);
      if (jumpback) BookmarkSet(jumpback, false);
      if (fardest) BpNew(fardest, true, false, 0, BPHARDWARE, 0, nullptr);
      DebugUpdateGuiAsync(addr, false);
      Sleep(300);

      // continue algorithm
      const unsigned char* curData = (addr >= _base && addr < _base + _size)
                                         ? _data() + (addr - _base)
                                         : nullptr;
      if (_cp.Disassemble(addr, curData, MAX_DISASM_BUFFER)) {
        if (addr + _cp.Size() >
            maxaddr)  // we went past the maximum allowed address
          break;

        if ((_cp.IsJump() || _cp.IsLoop()) && _cp.OpCount() &&
            _cp[0].type == ZYDIS_OPERAND_TYPE_IMMEDIATE)  // jump
        {
          duint dest = (duint)_cp[0].imm.value.u;

          if (dest >= maxaddr)  // jump across function boundaries
          {
            // currently unused
          } else if (dest > addr &&
                     dest >
                         fardest)  // save the farthest JXX destination forward
          {
            fardest = dest;
          } else if (end && dest < end &&
                     _cp.GetId() ==
                         ZYDIS_MNEMONIC_JMP)  // save the last JMP backwards
          {
            jumpback = addr;
          }
        } else if (_cp.IsRet())  // possible function end?
        {
          end = addr;
          if (fardest < addr)  // we stop if the farthest JXX destination
                               // forward is before this RET
            break;
        }

        addr += _cp.Size();
      } else
        addr++;
    }
    end = end < jumpback ? jumpback : end;

    // update GUI
    FunctionAdd(start, end, false);
    BpClear();
    BookmarkClear();
    SetContextDataEx(fdProcessInfo->hThread, UE_CIP, start);
    DebugUpdateGuiAsync(start, false);
  }
  return true;
}

bool cbInstrMeminfo(int argc, char* argv[]) {
  if (argc < 3) {
    dputs_untranslated("Usage: meminfo a/r, addr[, size]");
    return false;
  }
  duint addr;
  if (!valfromstring(argv[2], &addr)) {
    dputs_untranslated("Invalid argument");
    return false;
  }
  if (argv[1][0] == 'a') {
    duint size = 1;
    if (argc > 3 && !valfromstring(argv[3], &size)) {
      dputs_untranslated("Invalid argument");
      return false;
    }
    std::vector<uint8_t> buf;
    buf.resize(size);
    SIZE_T NumberOfBytesRead = 0;
    ReadProcessMemory(fdProcessInfo->hProcess, (const void*)addr, buf.data(),
                      buf.size(), &NumberOfBytesRead);
    dprintf_untranslated(
        "Data: %s\n",
        StringUtils::ToHex(buf.data(), NumberOfBytesRead).c_str());
  } else if (argv[1][0] == 'r') {
    MemUpdateMap();
    GuiUpdateMemoryView();
    dputs_untranslated("Memory map updated!");
  }
  return true;
}

bool cbInstrBriefcheck(int argc, char* argv[]) {
  if (IsArgumentsLessThan(argc, 2)) return false;
  duint addr;
  if (!valfromstring(argv[1], &addr, false)) return false;
  duint size;
  auto base = DbgMemFindBaseAddr(addr, &size);
  if (!base) return false;
  Memory<unsigned char*> buffer(size + 16);
  DbgMemRead(base, buffer(), size);
  Zydis cp;
  std::unordered_set<String> reported;
  for (duint i = 0; i < size;) {
    if (!cp.Disassemble(base + i, buffer() + i, 16)) {
      i++;
      continue;
    }
    i += cp.Size();
    auto mnem = StringUtils::ToLower(cp.Mnemonic());
    auto brief = MnemonicHelp::getBriefDescription(mnem.c_str());
    if (brief.length() || reported.count(mnem)) continue;
    reported.insert(mnem);
    dprintf_untranslated("%p: %s\n", cp.Address(), mnem.c_str());
  }
  return true;
}

bool cbInstrFocusinfo(int argc, char* argv[]) {
  ACTIVEVIEW activeView;
  GuiGetActiveView(&activeView);
  dprintf_untranslated("activeTitle: %s, activeClass: %s\n", activeView.title,
                       activeView.className);
  return true;
}

bool cbInstrFlushlog(int argc, char* argv[]) {
  GuiFlushLog();
  return true;
}

extern char animate_command[deflen];

bool cbInstrAnimateWait(int argc, char* argv[]) {
  while (DbgIsDebugging() && dbgisrunning() &&
         animate_command[0] != 0)  // while not locked (NOTE: possible deadlock)
  {
    Sleep(1);
  }
  return true;
}

#include <lz4/lz4file.h>

bool cbInstrDbdecompress(int argc, char* argv[]) {
  if (argc < 2) {
    dprintf_untranslated("Usage: dbdecompress \"c:\\path\\to\\db\"\n");
    return false;
  }
  auto dbFile = StringUtils::Utf8ToUtf16(argv[1]);
  if (LZ4_decompress_fileW(dbFile.c_str(), dbFile.c_str()) != LZ4_SUCCESS) {
    dprintf_untranslated("Failed to decompress '%s'\n", argv[1]);
    return false;
  }
  dprintf_untranslated("Decompressed '%s'\n", argv[1]);
  return true;
}

bool cbInstrDebugFlags(int argc, char* argv[]) {
  if (argc < 2) {
    dprintf_untranslated("Usage: DebugFlags 0xFFFFFFFF\n");
    return false;
  }
  auto debugFlags = DbgValFromString(argv[1]);
  dbgsetdebugflags(debugFlags);
  dprintf_untranslated("DebugFlags = 0x%08X\n", debugFlags);
  return true;
}

bool cbInstrLabelRuntimeFunctions(int argc, char* argv[]) {
#ifdef _WIN64
  if (argc < 2) {
    dputs_untranslated("Usage: LabelRuntimeFunctions modaddr");
    return false;
  }
  auto modaddr = DbgValFromString(argv[1]);
  SHARED_ACQUIRE(LockModules);
  auto info = ModInfoFromAddr(modaddr);
  if (info) {
    std::vector<COMMENTSINFO> comments;
    CommentGetList(comments);
    for (const auto& comment : comments) {
      if (comment.modhash == info->hash) {
        if (!comment.manual && comment.text.find("RUNTIME_FUNCTION") == 0) {
          CommentDelete(comment.addr + info->base);
        }
      }
    }
    for (const auto& runtimeFunction : info->runtimeFunctions) {
      auto setComment = [info](duint addr, const char* prefix) {
        char comment[MAX_COMMENT_SIZE] = "";
        if (!CommentGet(addr, comment))
          strncpy_s(comment, "RUNTIME_FUNCTION", _TRUNCATE);
        strncat_s(comment, " ", _TRUNCATE);
        strncat_s(comment, prefix, _TRUNCATE);
        CommentSet(addr, comment, false);
      };
      setComment(info->base + runtimeFunction.BeginAddress, "BeginAddress");
      setComment(info->base + runtimeFunction.EndAddress, "EndAddress");
    }
    GuiUpdateAllViews();
  } else {
    dprintf_untranslated("No module found at %p\n", modaddr);
  }
  return true;
#else
  return false;
#endif  // _WIN64
}
